close all; clear; clc;

% 3D FDTD simulation with PLM
% FDTD main loop seems to be Ok but source inject has problems
%% Common value for simulation
c0 = 299792458; % light speed
e0 = 8.8541878176e-12;% permittivity of free space
u0 = 1.2566370614e-6; % permeability of free space
eta0 = sqrt(u0/e0); % impedance of free space

gigahertz = 1e9;
megahertz = 1e6;

ER_air = 1;
UR_air = 1;
nmax = sqrt(ER_air*UR_air); % refractive index for material, here is air

freq_max = 1*gigahertz; % 1GHz

x_length = 1.5; % 1m in x axis
y_length = 1.5; % 1m in y axis
z_length = 2; % 2m in z axis


%% Start calculate grid resolution
NRES = 20; % resolve the wave with at least 10 cells, better >=10
lambda_min = c0./freq_max./nmax;
delta_lambda = lambda_min/NRES;

% use free space in this example so material resultion is not needed
%NDRES = 10; % normally 1~4 for resolution of feature size
%delta_d = x_length/NDRES;
%d_temp = min(delta_lambda, delta_d);

d_temp = delta_lambda;

Nx = ceil(x_length/d_temp);
Ny = ceil(y_length/d_temp);
Nz = ceil(z_length/d_temp);

dx = x_length/Nx; % get the size of X grid
dy = y_length/Ny; % get the size of Y grid
dz = z_length/Nz; % get the size of Z grid
d_avg = (dx+dy+dz)/3;

% cells for NPML:NXLO;NXHI;NYLO;NYHI;NZLO;NZHI
NPML = [15;15;15;15;15;15]; % cells for PML in 1x Grid

%% Start calculate delta_t and tau
Nt = 10; % resuolution for t, normally >=10;
dt1 = 1/Nt/freq_max; 
dt2 = min(dx,dy)/(2*c0); % also consider 1 cell 2 time steps at least
dt = min(dt1,dt2);
delta_freq = 1*megahertz; % target frequency resolution, here is 1Mhz

% STEPS>=1/(dt*delta_f) ths is to make sure freq resoltuion in freq domain
% STEPS = 1/dt/delta_freq;

% Signal as Gauss pusle width if using Gauss Source
tau = 1/(2*freq_max);
t0 = 6*tau;

% Signal as single tone
% tau = 3/freq_max; % sin source
% t0 = 3*tau;

% another way to cal STEPS is to calculate prop time if freq domain does
% not need to be considered
t_prop = max(max(nmax*Nx*dx/c0, nmax*Ny*dy/c0), nmax*Nz*dz/c0);
T_total = 12*tau + 5*t_prop;
STEPS = ceil(T_total/dt); % STEPS calculated by prop time

t_axis=(0:STEPS-1).*dt; % time axis

%% 2x Grid calculation
Nx2 = 2*Nx; Ny2 = 2*Ny; Nz2 = 2*Nz;
dx2 = dx/2; dy2 = dy/2; dz2 = dz/2;

% calculate ER/UR based on 2x grid
% this example is free space and always 1

ER2 = ones(Nx2,Ny2,Nz2); % isotropic for free space
UR2 = ones(Nx2,Ny2,Nz2); % isotropic for free space

% after ER2 and ER2 ready, notmally need to 'smooth' it  by convolution
% this example is free space and no need smooth

% calculate for 3D scenario
ERxx = ER2(2:2:Nx2,1:2:Ny2,1:2:Nz2);
ERyy = ER2(1:2:Nx2,2:2:Ny2,1:2:Nz2);
ERzz = ER2(1:2:Nx2,1:2:Ny2,2:2:Nz2);
URxx = UR2(1:2:Nx2,2:2:Ny2,2:2:Nz2);
URyy = UR2(2:2:Nx2,1:2:Ny2,2:2:Nz2);
URzz = UR2(2:2:Nx2,2:2:Ny2,1:2:Nz2);

%% Gause Source definition
Nx_src = 50;
Ny_src = 50;
Nz_src = 50;
ER_src = ERxx(Nx_src, Ny_src, Nz_src); % did not consider anisotropy
UR_src = URxx(Nx_src, Ny_src, Nz_src); % did not consider anisotropy
n_src = sqrt(ER_src*UR_src);

delta_t = n_src*dz/(2*c0) + dt/2; % total delay between E and H
Esrc = exp(-((t_axis-t0)/tau).^2); % E filed source
% A = -sqrt(ER_src/UR_src); % amplitude of H field
% Hsrc = A*exp(-((t_axis-t0+delta_t)/tau).^2); % H field source

%% Calculate PML parameters
sigx = zeros(Nx2,Ny2,Nz2);
for nx = 1 : 2*NPML(1)
    nx1 = 2*NPML(1) - nx + 1;
    sigx(nx1,:,:) = (0.5*e0/dt)*(nx/2/NPML(1))^3;
end
for nx = 1 : 2*NPML(2)
    nx1 = Nx2 - 2*NPML(2) + nx;
    sigx(nx1,:,:) = (0.5*e0/dt)*(nx/2/NPML(2))^3;
end

sigy = zeros(Nx2,Ny2,Nz2);
for ny = 1 : 2*NPML(3)
    ny1 = 2*NPML(3) - ny + 1;
    sigy(:,ny1,:) = (0.5*e0/dt)*(ny/2/NPML(3))^3;
end
for ny = 1 : 2*NPML(4)
    ny1 = Ny2 - 2*NPML(4) + ny;
    sigy(:,ny1,:) = (0.5*e0/dt)*(ny/2/NPML(4))^3;
end

sigz = zeros(Nx2,Ny2,Nz2);
for nz = 1 : 2*NPML(5)
    nz1 = 2*NPML(5) - nz + 1;
    sigz(:,:,nz1) = (0.5*e0/dt)*(nz/2/NPML(5))^3;
end
for nz = 1 : 2*NPML(6)
    nz1 = Nz2 - 2*NPML(6) + nz;
    sigz(:,:,nz1) = (0.5*e0/dt)*(nz/2/NPML(6))^3;
end

%% fetch PML function onto 1x
% calculate updated coefficients for 3D model
sigHx = sigx(1:2:Nx2,2:2:Ny2,2:2:Nz2);
sigHy = sigy(1:2:Nx2,2:2:Ny2,2:2:Nz2);
sigHz = sigz(1:2:Nx2,2:2:Ny2,2:2:Nz2);
mHx0 = (1/dt) + (sigHy + sigHz)/(2*e0) ...
    + sigHy.*sigHz*(dt/4/e0^2);
mHx1 = (1/dt) - (sigHy + sigHz)/(2*e0) ...
    - sigHy.*sigHz*(dt/4/e0^2);
mHx1 = mHx1./mHx0;
mHx2 = - c0./URxx./mHx0;
mHx3 = - (c0*dt/e0) * sigHx./URxx./mHx0;
mHx4 = - (dt/e0^2)*sigHy.*sigHz./mHx0;

sigHx = sigx(2:2:Nx2,1:2:Ny2,2:2:Nz2);
sigHy = sigy(2:2:Nx2,1:2:Ny2,2:2:Nz2);
sigHz = sigy(2:2:Nx2,1:2:Ny2,2:2:Nz2);
mHy0 = (1/dt) + (sigHx + sigHz)/(2*e0) ...
    + sigHx.*sigHz*(dt/4/e0^2);
mHy1 = (1/dt) - (sigHx + sigHz)/(2*e0) ...
    - sigHx.*sigHz*(dt/4/e0^2);
mHy1 = mHy1./mHy0;
mHy2 = - c0./URyy./mHy0;
mHy3 = - (c0*dt/e0) * sigHy./URyy./mHy0;
mHy4 = - (dt/e0^2)*sigHx.*sigHz./mHy0;

sigHx = sigx(2:2:Nx2,2:2:Ny2,1:2:Nz2);
sigHy = sigy(2:2:Nx2,2:2:Ny2,1:2:Nz2);
sigHz = sigy(2:2:Nx2,2:2:Ny2,1:2:Nz2);
mHz0 = (1/dt) + (sigHx + sigHy)/(2*e0) ...
    + sigHx.*sigHy*(dt/4/e0^2);
mHz1 = (1/dt) - (sigHx + sigHy)/(2*e0) ...
    - sigHx.*sigHy*(dt/4/e0^2);
mHz1 = mHz1./mHz0;
mHz2 = - c0./URzz./mHz0;
mHz3 = - (c0*dt/e0) * sigHz./URzz./mHz0;
mHz4 = - (dt/e0^2)*sigHx.*sigHy./mHz0;

sigDx = sigx(1:2:Nx2,1:2:Ny2,1:2:Nz2);
sigDy = sigy(1:2:Nx2,1:2:Ny2,1:2:Nz2);
sigDz = sigz(1:2:Nx2,1:2:Ny2,1:2:Nz2);
mDx0 = (1/dt) + (sigDy + sigDz)/(2*e0) ...
    + sigDy.*sigDz*(dt/4/e0^2);
mDx1 = (1/dt) - (sigDy + sigDz)/(2*e0) ...
    - sigDy.*sigDz*(dt/4/e0^2);
mDx1 = mDx1 ./ mDx0;
mDx2 = c0./mDx0;
mDx3 = (c0*dt/e0) * sigDx./mDx0;
mDx4 = - (dt/e0^2)*sigDy.*sigDz./mDx0;
mEx1 = 1./ERxx;

sigDx = sigx(1:2:Nx2,1:2:Ny2,1:2:Nz2);
sigDy = sigy(1:2:Nx2,1:2:Ny2,1:2:Nz2);
sigDz = sigz(1:2:Nx2,1:2:Ny2,1:2:Nz2);
mDy0 = (1/dt) + (sigDx + sigDz)/(2*e0) ...
    + sigDx.*sigDz*(dt/4/e0^2);
mDy1 = (1/dt) - (sigDx + sigDz)/(2*e0) ...
    - sigDx.*sigDz*(dt/4/e0^2);
mDy1 = mDy1 ./ mDy0;
mDy2 = c0./mDy0;
mDy3 = (c0*dt/e0) * sigDy./mDy0;
mDy4 = - (dt/e0^2)*sigDx.*sigDz./mDy0;
mEy1 = 1./ERyy;

sigDx = sigx(1:2:Nx2,1:2:Ny2,1:2:Nz2);
sigDy = sigy(1:2:Nx2,1:2:Ny2,1:2:Nz2);
sigDz = sigz(1:2:Nx2,1:2:Ny2,1:2:Nz2);
mDz0 = (1/dt) + (sigDx + sigDy)/(2*e0) ...
    + sigDx.*sigDy*(dt/4/e0^2);
mDz1 = (1/dt) - (sigDx + sigDy)/(2*e0) ...
    - sigDx.*sigDy*(dt/4/e0^2);
mDz1 = mDz1 ./ mDz0;
mDz2 = c0./mDz0;
mDz3 = (c0*dt/e0) * sigDz./mDz0;
mDz4 = - (dt/e0^2)*sigDx.*sigDy./mDz0;
mEz1 = 1./ERzz;


%% Initialize field to zero
Hx = zeros(Nx,Ny,Nz); Hy = zeros(Nx,Ny,Nz); Hz = zeros(Nx,Ny,Nz);
Ex = zeros(Nx,Ny,Nz); Ey = zeros(Nx,Ny,Nz); Ez = zeros(Nx,Ny,Nz);
Dx = ERxx.*Ex; Dy = ERyy.*Ey; Dz = ERzz.*Ez;
CEx = zeros(Nx,Ny,Nz); CEy = zeros(Nx,Ny,Nz); CEz = zeros(Nx,Ny,Nz);
CHx = zeros(Nx,Ny,Nz); CHy = zeros(Nx,Ny,Nz); CHz = zeros(Nx,Ny,Nz);
ICEx = zeros(Nx,Ny,Nz); ICEy = zeros(Nx,Ny,Nz); ICEz = zeros(Nx,Ny,Nz);
ICHx = zeros(Nx,Ny,Nz); ICHy = zeros(Nx,Ny,Nz); ICHz = zeros(Nx,Ny,Nz);
IDx = zeros(Nx,Ny,Nz); IDy = zeros(Nx,Ny,Nz); IDz = zeros(Nx,Ny,Nz);
IHx = zeros(Nx,Ny,Nz); IHy = zeros(Nx,Ny,Nz); IHz = zeros(Nx,Ny,Nz);

%% initialize plot parameters
fig=figure;
set(fig,'Name', 'FDTD 3D PML Simulation');
set(fig,'NumberTitle', 'off');
xa=1:Nx;
ya=1:Ny;
za=1:Nz;

[Y_xoy,X_xoy] = meshgrid(ya,xa);
[Z_xoz,X_xoz] = meshgrid(za,xa);
[Z_yoz,Y_yoz] = meshgrid(za,ya);
[X,Y,Z] = meshgrid(xa,ya,za);

%% Core FDTD Loop for 3D
for T = 1 : STEPS
    % Calculate curl of E, update CEx
    for nx = 1 : Nx
        for ny = 1 : Ny-1
            for nz = 1 : Nz-1
                CEx(nx,ny,nz) = (Ez(nx,ny+1,nz) - Ez(nx,ny,nz))/dy ...
                    - (Ey(nx,ny,nz+1) - Ey(nx,ny,nz))/dz;
            end
            CEx(nx,ny,Nz) = (Ez(nx,ny+1,Nz) - Ez(nx,ny,Nz))/dy ...
                - (0 - Ey(nx,ny,Nz))/dz;
        end
        for nz = 1 : Nz-1
            CEx(nx,Ny,nz) = (0 - Ez(nx,Ny,nz))/dy ...
                - (Ey(nx,Ny,nz+1) - Ey(nx,Ny,nz))/dz;
        end
        CEx(nx,Ny,Nz) = (0 - Ez(nx,Ny,Nz))/dy ...
            - (0 - Ey(nx,Ny,Nz))/dz;
    end
    % Calculate curl of E, update CEy
    for ny = 1 : Ny
        for nz = 1 : Nz-1
            for nx = 1 : Nx-1
                CEy(nx,ny,nz) = (Ex(nx,ny,nz+1) - Ex(nx,ny,nz))/dz ...
                    - (Ez(nx+1,ny,nz) - Ez(nx,ny,nz))/dx;
            end
            CEy(Nx,ny,nz) = (Ex(Nx,ny,nz+1) - Ex(Nx,ny,nz))/dz ...
                - (0 - Ez(Nx,ny,nz))/dx;
        end
        for nx = 1 : Nx-1
            CEy(nx,ny,Nz) = (0 - Ex(nx,ny,Nz))/dz ...
                - (Ez(nx+1,ny,Nz) - Ez(nx,ny,Nz))/dx;
        end
        CEy(Nx,ny,Nz) = (0 - Ex(Nx,ny,Nz))/dz ...
            - (0 - Ez(Nx,ny,Nz))/dx;
    end
    % Calculate curl of E, update CEz
    for nz = 1 : Nz
        for nx = 1 : Nx-1
            for ny = 1 : Ny-1
                CEz(nx,ny,nz) = (Ey(nx+1,ny,nz) - Ey(nx,ny,nz))/dx ...
                    - (Ex(nx,ny+1,nz) - Ex(nx,ny,nz))/dy;
            end
            CEz(nx,Ny,nz) = (Ey(nx+1,Ny,nz) - Ey(nx,Ny,nz))/dx ...
                - (0 - Ex(nx,Ny,nz))/dy;
        end
        for ny = 1 : Ny-1
            CEz(Nx,ny,nz) = (0 - Ey(Nx,ny,nz))/dx ...
                - (Ex(Nx,ny+1,nz) - Ex(Nx,ny,nz))/dy;
        end
        CEz(Nx,Ny,nz) = (0 - Ey(Nx,Ny,nz))/dx ...
            - (0 - Ex(Nx,Ny,nz))/dy;
    end

    % Update H Integrations
    ICEx = ICEx + CEx; ICEy = ICEy + CEy; ICEz = ICEz + CEz;
    IHx = IHx + Hx; IHy = IHy + Hy; IHz = IHz + Hz;
    
    % Update H from E
    Hx = mHx1.*Hx + mHx2.*CEx + mHx3.*ICEx + mHx4.*IHx;
    Hy = mHy1.*Hy + mHy2.*CEy + mHy3.*ICEy + mHy4.*IHy;
    Hz = mHz1.*Hz + mHz2.*CEz + mHz3.*ICEz + mHz4.*IHz;
     
    % Calculate curl of H, update CHx
    for nx = 1 : Nx
        CHx(nx,1,1) = (Hz(nx,1,1) - 0)/dy ...
            - (Hy(nx,1,1) - 0)/dz;
        for nz = 2 : Nz
            CHx(nx,1,nz) = (Hz(nx,1,nz) - 0)/dy ...
                - (Hy(nx,1,nz) - Hy(nx,1,nz-1))/dz;
        end
        for ny = 2 : Ny
            CHx(nx,ny,1) = (Hz(nx,ny,1) - Hz(nx,ny-1,1))/dy ...
                - (Hy(nx,ny,1) - 0)/dz;
            for nz = 2 : Nz
                CHx(nx,ny,nz) = (Hz(nx,ny,nz) - Hz(nx,ny-1,nz))/dy ...
                    - (Hy(nx,ny,nz) - Hy(nx,ny,nz-1))/dz;
            end
        end
    end
    % Calculate curl of H, update CHy
    for ny  = 1 : Ny
        CHy(1,ny,1) = (Hx(1,ny,1) - 0)/dz ...
            - (Hz(1,ny,1) - 0)/dx;
        for nx = 2 : Nx
            CHy(nx,ny,1) = (Hx(nx,ny,1) - 0)/dz ...
                - (Hz(nx,ny,1) - Hz(nx-1,ny,1))/dx;
        end
        for nz = 2 : Nz
            CHy(1,ny,nz) = (Hx(1,ny,nz) - Hx(1,ny,nz-1))/dz ...
                - (Hz(1,ny,nz) - 0)/dx;
            for nx = 2 : Nx
                CHy(nx,ny,nz) = (Hx(nx,ny,nz) - Hx(nx,ny,nz-1))/dz ...
                    - (Hz(nx,ny,nz) - Hz(nx-1,ny,nz))/dx;
            end
        end
    end
    % Calculate curl of H, update CHz
    for nz = 1 : Nz
        CHz(1,1,nz) = (Hy(1,1,nz) - 0)/dx ...
            - (Hx(1,1,nz) - 0)/dy;
        for ny = 2 : Ny
            CHz(1,ny,nz) = (Hy(1,ny,nz) - 0)/dx ...
                - (Hx(1,ny,nz) - Hx(1,ny-1,nz))/dy;
        end
        for nx = 2 : Nx
            CHz(nx,1,nz) = (Hy(nx,1,nz) - Hy(nx-1,1,nz))/dx ...
                - (Hx(nx,1,nz) - 0)/dy;
            for ny = 2 : Ny
                CHz(nx,ny,nz) = (Hy(nx,ny,nz) - Hy(nx-1,ny,nz))/dx ...
                    - (Hx(nx,ny,nz) - Hx(nx,ny-1,nz))/dy;
            end
        end
    end
        
    % Update D Integrations
    IDx = IDx + Dx; IDy = IDy + Dy; IDz = IDz + Dz;
    % ICH happened in t-dt/2 or t+dt/2 ???, here need to be checked
    ICHx = ICHx + CHx; ICHy = ICHy + CHy; ICHz = ICHz + CHz;
    
    % Update Dz
    Dx = mDx1.*Dx + mDx2.*CHx + mDx3.*ICHx + mDx4.*IDx;
    Dy = mDy1.*Dy + mDy2.*CHy + mDy3.*ICHy + mDy4.*IDy;
    Dz = mDz1.*Dz + mDz2.*CHz + mDz3.*ICHz + mDz4.*IDz;
    
    % Inject Source
    Dx(Nx_src,Ny_src,Nz_src) = Esrc(T) * ERzz(Nx_src,Ny_src,Nz_src);
    Dy(Nx_src,Ny_src,Nz_src) = Esrc(T) * ERzz(Nx_src,Ny_src,Nz_src);
    Dz(Nx_src,Ny_src,Nz_src) = Esrc(T) * ERzz(Nx_src,Ny_src,Nz_src);
    
    % Update Ez
    Ex = mEx1.*Dx;
    Ey = mEy1.*Dy;
    Ez = mEz1.*Dz;
    
    
    %Plot
    subplot(1,2,1);
    E_power = abs(permute(Ez,[1 2 3])).^2 ...
        + abs(permute(Ex,[1 2 3])).^2 ...
        + abs(permute(Ey,[1 2 3])).^2;
    E_power_log = 10*log10(E_power);
    
    h1=slice(X,Y,Z,E_power_log, [Nx_src,Nx_src+20], [], Nz_src);
    caxis([-100,0]);
    set(h1,'edgecolor','none'); alpha(h1,0.6); axis equal tight; colormap jet;
    xlabel('OX, m'); ylabel('OY, m'); zlabel('OZ, m'); axis equal tight;
    title(sprintf('Step: %d of %d',T, STEPS));
    
    
    subplot(1,2,2);
    E_power_2 = abs(squeeze(Ez(:,Ny_src,:))).^2 ...
        + abs(squeeze(Ex(:,Ny_src,:))).^2 ...
        + abs(squeeze(Ey(:,Ny_src,:))).^2;
    
    pcolor(X_xoz,Z_xoz,10*log10(E_power_2));
    axis ij;
    axis equal tight;
    colormap('jet');
    shading interp;
    caxis([-100,0]);
    xlabel('OX, m'); ylabel('OZ, m');
    %colorbar
%     E_power = abs(squeeze(Ex(:,:,Nz_src))).^2 ...
%         + abs(squeeze(Ey(:,:,Nz_src))).^2 ...
%         + abs(squeeze(Ez(:,:,Nz_src))).^2;
%     pcolor(X_xoy,Y_xoy,10*log10(E_power));
%     axis ij
%     axis equal tight;
%     colormap('jet');
%     shading interp;
%     caxis([-100,0]);
%     %colorbar
%     hold on
    title(sprintf('Step: %d of %d',T, STEPS));
    
    drawnow;
    hold off
    %pause(0.01);
    
end

